import re
import collections
import tokenize
import operator
import itertools

# ___________________________________________________
#|:                                            ``::%H|
#|%:.  a function for transversing JSON objects   `:%|   
#|H%::..____________________________________________:|

def jpath(obj, path):
    """
    Transverse 'obj' to get the content in the given 'path'
    
    Arguments
    ---------
    
    obj : JSON-like object 
        JSON structure (list-like or dictionary-like)
    
    path : str
        JPATH location of object
    
    
    Examples
    --------------
    
    Given the JSON-like structure
    
    { "store": {
        "book": [ 
          { "category": "reference",
            "author": "Nigel Rees",
            "title": "Sayings of the Century",
            "price": 8.95
          },
          { "category": "fiction",
            "author": "Evelyn Waugh",
            "title": "Sword of Honour",
            "price": 12.99
          },
          { "category": "fiction",
            "author": "Herman Melville",
            "title": "Moby Dick",
            "isbn": "0-553-21311-3",
            "price": 8.99
          },
          { "category": "fiction",
            "author": "J. R. R. Tolkien",
            "title": "The Lord of the Rings",
            "isbn": "0-395-19395-8",
            "price": 22.99
          }
        ],
        "bicycle": {
          "color": "red",
          "price": 19.95
        }
      }
    }
    
    JPATH can be used to make the following queries
    
        $.store.book[*].author  --- the authors of all books in the store
        $.store..price          --- the prices of everything in the store.
        $.store.book[2]         --- the third book
        $..book[-1]             --- the last book (recursive descent)
        $.store.book            --- ordered list of books
        $.store.book[::-1]      --- list of books in reversed order
        $.store.book[0,1,4]     --- first, second and fifth books
        $.store.book[?(isbn)]   --- filter all books with isbn number
        $..book[?(price < 10)]  --- filter all books cheaper than 10
        $..*                    --- All members of JSON structure.
    
    Observations
    ------------
    
    JPATH is vaguely based on the notation defined at
    http://goessner.net/articles/JsonPath/.
    """

    return JPathWorker(path)(obj, path[1:])

class JPathWorker(object):
    def __init__(self, path):
        if not path.startswith('$.'):
            raise ValueError("invalid path '%s', must start with '$.'" % path)
        self.orig_path = path

    def walk(self, obj, path):
        m = re_walk.match(path)
        while m:
            key = m.groups()[0]
            try:
                obj = obj[key]
            except TypeError:
                obj = obj[int(key)]
            path = path[m.end():].lstrip()
            if path.startswith('['):
                op, _, path = path[1:].partition(']')
                obj = self.array_operation(obj, op)
            m = re_walk.match(path)

        return obj, path

    def array_operation(self, obj, array_op):
        array_op = array_op.replace(' ', '')

        # filter
        if array_op.startswith('?(') and array_op.endswith(')'):
            filter = array_op[2:-1]
            return self.array_filter(obj, filter)

        # sort
        if array_op.startswith('#(') and array_op.endswith(')'):
            key = array_op[2:-1]
            return self.array_sort(obj, key)

        # select by list of indexes
        m = re_cvs.match(array_op)
        if m:
            idxs = map(int, array_op.split(','))
            return [ obj[idx] for idx in idxs ]

        # slice operations
        if ':' in array_op:
            slice = [ (int(x) if x else None) for x in array_op.split(':') ]
            if len(slice) == 2:
                a, b = slice
                return obj[a:b]
            elif len(slice) == 3:
                a, b, c = slice
                return obj[a:b:c]
            else:
                raise IndexError("Invalid slice '%s' at %s" % (array_op, self.orig_path))

        raise IndexError("unsupported array index '%s' in path '%s'" % (array_op, self.orig_path))

    def array_filter(self, obj, filter_):
        """filter array given by obj"""
        try:
            key, op, num, st = re_filter.findall(filter_)[0]
        except IndexError:
            raise KeyError("invalid array filter: ?(%s)" % filter_)
        keys = key.split('.')

        # define the getter function
        key = key.lstrip('@')
        walk = self.walk
        def getitem(x):
            try:
                obj, p = walk(x, key)
                assert not p, 'invalid traling path: %s' % p
                return obj
            except KeyError:
                return None

        # simple filter: check if properties exist
        if not op:
            f = lambda x: not (getitem(x) is None)
            return filter(f, obj)

        # define acceptance criteria
        op = filter_operators[op]
        if num:
            try:
                other = int(num)
            except ValueError:
                other = float(num)
        elif st:
            other = st[1:-1].decode('string_escape')
        else:
            raise KeyError("invalid filter '%s'" % filter)

        # check which elements in the array satisfy the acceptance criteria
        def accept(x):
            item = getitem(x)
            if not other is None:
                return op(item, other)
        return filter(accept, obj)

    def array_sort(self, obj, key):
        """filter array given by obj"""
        keys = key.split('.')

        # define the getter function
        key = key.lstrip('@')
        walk = self.walk
        def getitem(x):
            try:
                obj, p = walk(x, key)
                assert not p, 'invalid traling path: %s' % p
                return obj
            except KeyError:
                return None

        # simple filter: check if properties exist
        items = [ (x, getitem(x)) for x in obj ]
        items.sort(key=lambda x: x[1])
        nulls = [ x[0] for x in items if x[1] is None ]
        non_nulls = [ x[0] for x in items if not x[1] is None ]
        return non_nulls + nulls

    def recursive_descent(self, obj, path):
        visited = self.visited
        qset = self.queryset
        if id(obj) in visited:
            return
        else:
            visited.append(id(obj))

        if isinstance(obj, (basestring, int, float, bool, type(None))):
            return

        if isinstance(obj, collections.Mapping):
            try:
                self._rdescent_item(obj, path)
            except:
                pass
            obj = obj.values()

        if isinstance(obj, collections.Sequence):
            for item in obj:
                try:
                    self._rdescent_item(obj, path)
                except:
                    pass
                self.recursive_descent(item, path)

    def _rdescent_item(self, obj, path):
        item = self(obj, path)
        if not id(item) in self.visited:
            self.visited.append(id(item))
            if isinstance(item, queryset):
                self.queryset.extend(item)
            else:
                self.queryset.append(item)

    def expand_asterisk(self, obj, path):
        if isinstance(obj, collections.Mapping):
            seq = obj.items()
        else:
            seq = iter(obj)

        expansions = (self(o, path) for o in obj)
        result = queryset()
        for item in obj:
            try:
                item = self(item, path)
            except (KeyError, IndexError):
                continue
            if isinstance(item, queryset):
                result.extend(item)
            else:
                result.append(item)
        return result

    def __call__(self, obj, path):
        path = path.strip()
        if not path:
            return obj
        elif path.startswith('.*'):
            path = path[2:]
            return self.expand_asterisk(obj, path)
        elif path.startswith('..'):
            path = path[1:]
            self.visited = set()
            self.queryset = []
            self.recursive_descent(obj, path)
            return self.queryset
        else:
            # simple walk to object
            obj, path = self.walk(obj, path)
            if path:
                return self(obj, path)
            else:
                return obj

#===============================================================================
#                   Helper functions: REGEX and constants
#===============================================================================
class queryset(list): pass

re_cvs = re.compile('^\d+\s*(,\d+\s*)*\s*\Z')
re_sort = re.compile(r'^\s*#.\s*([a-zA-Z]\w*)')
re_walk = re.compile('^[.]([a-zA-Z]\w*|[-]?[0-9]+)')
re_filter = re.compile(r'^\s*(@.(?:[a-zA-Z]\w*))\s*(==|!=|>=|<=|<|>|is)?\s*(%s)?%s?\s*\Z' %
                    (tokenize.Number.replace('[lL]+', '').replace('(', '(?:'),
                     tokenize.String))

filter_operators = {
    '==': operator.eq,
    '!=': operator.ne,
    '>=': operator.ge,
    '<=': operator.le,
    '>': operator.gt,
    '<': operator.lt, }

#===============================================================================
#                                TEST
#===============================================================================
if __name__ == '__main__':
    qtest = { 'meta': { 'title': 'Foo', 'db_name': 'foobar' },
              'content': [ { 'type': 'text',
                             'text': 'some text' },
                           { 'type': 'multiple-choice',
                             'items': [ { 'text': 'Foo' },
                                        { 'text': 'Bar' } ] }]}

    def _jpath(x, v):
        return '%s: %s' % (v, jpath(x, v))
    print _jpath(qtest, '$.meta')
    print _jpath(qtest, '$.meta.title')
    print _jpath(qtest, '$.content.0.type')
    print _jpath(qtest, '$.content.-1.type')
    print _jpath(qtest, '$.content')
    print _jpath(qtest, '$.content[0,1]')
    print _jpath(qtest, '$.content[::-1]')
    print _jpath(qtest, '$.content[0]')
    print _jpath(qtest, '$.content[?(@.text)]')
    print _jpath(qtest, '$.content[?(@.type == "text")]')
    print _jpath(qtest, "$.content[?(@.text == 'text')]")
    print _jpath(qtest, "$.content[?(@.text > 1)]")
    print _jpath(qtest, "$.content[?(@.text > 1)]")
    print _jpath(qtest, '$.content[#(@.text)]')
    print _jpath(qtest, '$.meta.*')
    print _jpath(qtest, '$.content[0,1].*.type')
    print _jpath(qtest, '$.content.*.type')
    print _jpath(qtest, '$.content.*.items.*.text')
    print _jpath(qtest, '$.content..text')
    print _jpath(qtest, '$.content..items.*.text')
#    print _jpath(qtest, "$.content[?(@.text > @.type)]")
#    print _jpath(qtest, '$.content[*].items[*].text')
#    print _jpath(qtest, '$.content[0,3]')
#    print _jpath(qtest, '$.content.3.type')
